from easy_io import read_mat_file, write_pkl_file
import numpy as np
import sys
sys.path.insert(0, '/data_4t/Kaggle/python_code/data')
from Reader import Label_Reader


def prepare_labels(candidate_mat_file):
    candidate_arr = read_mat_file(candidate_mat_file)
    candidate_lst = []
    for s in candidate_arr:
        scanid = str(s[0][0])
        assert s[1].ndim == 2
        for c in s[1]:
            c = np.asarray(c[:4], 'int')
            candidate_lst.append((scanid, c[:3], c[3]))

    reader = Label_Reader()
    last_scanid = ''
    labels = []
    for i, (scanid, midpoint, diameter) in enumerate(candidate_lst):
        # midpoint = np.array([midpoint[1], midpoint[0], midpoint[2]], 'int')
        try:
            if scanid != last_scanid:
                reader.bind(scanid)
            last_scanid = scanid
            labels.append(-1 if reader.eval_center([midpoint])[1][0] == np.True_ else 1)
        except:
            print(scanid)
            labels.append(0)
    return labels


if __name__ == '__main__':
    write_pkl_file(
        '/ssd_1t/Kaggle/lidc_candidate_label.pkl',
        prepare_labels('/data_4t/Kaggle/candidates/lidc_candis.mat')
    )
